package com.demo.websocket.interceptor;

import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.http.server.ServletServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.Map;

/**
 * @Author JiangMing Zhou
 * @Date 2023/4/20 16:42
 * @Version 1.8
 */
@Component
public class WebSocketInterceptor implements HandshakeInterceptor {

    /**
     * 握手前
     * @param serverHttpRequest
     * @param serverHttpResponse
     * @param webSocketHandler
     * @param map
     * @return
     * @throws Exception
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse, WebSocketHandler webSocketHandler, Map<String, Object> map) throws Exception {
        ServletServerHttpRequest request = (ServletServerHttpRequest) serverHttpRequest;
        ServletServerHttpResponse response = (ServletServerHttpResponse) serverHttpResponse;
        String token = request.getServletRequest().getHeader("Sec-WebSocket-Protocol");
        if (StringUtils.hasLength(token)){
            //token合法
            response.getServletResponse().setHeader("Sec-WebSocket-Protocol", token);
            map.put("uid","111");
            return true;
        }else {
            return false;
        }
    }

    /**
     * 握手后
     * @param serverHttpRequest
     * @param serverHttpResponse
     * @param webSocketHandler
     * @param e
     */
    @Override
    public void afterHandshake(ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse, WebSocketHandler webSocketHandler, Exception e) {

    }
}
